/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Description: Stub header generator.
 */
#include "datasystem/common/rpc/plugin_generator/zmq_rpc_generator.h"

namespace datasystem {

void ZmqRpcGenerator::CreateStubHeader(const google::protobuf::FileDescriptor &file,
                                       compiler::GeneratorContext *generatorCtx) const
{
    std::unique_ptr<io::ZeroCopyOutputStream> outputFile(generatorCtx->Open(fileName + ".stub.rpc.pb.h"));
    io::Printer printer(outputFile.get(), '$');

    GenerateStubPrologue(printer, file);

    printer.PrintRaw(namespaceBegin);

    for (auto i = 0; i < file.service_count(); ++i) {
        auto *svc = file.service(i);
        GenerateStubClass(printer, *svc, PREFIX);
        GenerateGenericStubClass(printer, *svc, PREFIX);
    }

    printer.PrintRaw(namespaceEnd);
    printer.PrintRaw(ENDIF);
}

void ZmqRpcGenerator::GenerateGenericStubClass(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc,
                                               const std::string &indent) const
{
    const std::string &svcName = svc.name();
    std::map<std::string, std::string> vars;
    vars["svc_name"] = svcName;
    vars["stub"] = svcName + "_Stub";
    vars["multi_session"] = std::to_string(MultiSessionEnabled(svc));
    const std::string constructor =
        "class $stub$ : public RpcStubBase {\n"
        "public:\n"
        "    explicit $stub$(std::shared_ptr<::datasystem::RpcChannel> channel, int32_t timeoutMs = -1);\n"
        "    ~$stub$() = default;\n"
        "    Status GetInitStatus() override;\n";
    printer.Print(vars, constructor.c_str());

    // Implement the override function ServiceName.
    GenerateSvcName(printer, svcName, indent, false);

    // Implement stub api.
    ImplementStubApiDecl(printer, svc, indent);

    // Implement adapter for some other api in stub
    ImplementGenericStubOtherFuncDecl(printer);

    const std::string impl =
        "private:\n"
        "    std::unique_ptr<$svc_name$_ZmqStub> stub_;\n"
        "};\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::GenerateStubClass(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc,
                                        const std::string &indent) const
{
    const std::string &svcName = svc.name();
    std::map<std::string, std::string> vars;
    vars["svc_name_string"] = "\"" + svcName + "\"";
    vars["stub"] = svcName + "_ZmqStub";
    vars["channel_number"] = std::to_string(HasChannelOption(svc));
    const std::string constructor =
        "class $stub$ final : public ::datasystem::ZmqStub {\n"
        "public:\n"
        "    explicit $stub$ (std::shared_ptr<::datasystem::RpcChannel> channel, int32_t timeoutMs = -1)\n"
        "    : ::datasystem::ZmqStub(std::move(channel), timeoutMs) \n"
        "    { InitMethodMap(); serviceName_ = $svc_name_string$; \n"
        "      channelNo_ = $channel_number$; \n"
        "      this->InitConn(); }\n"
        "    ~$stub$() override = default;\n";
    printer.Print(vars, constructor.c_str());

    // Create the InitMethodMap.
    GenerateInitMethodMapDecl(printer);

    // Implement the override function ServiceName.
    GenerateSvcName(printer, svcName, indent);

    // Implement stub api.
    ImplementStubApiDecl(printer, svc, indent);

    printer.PrintRaw("private:\n");
    printer.PrintRaw("};\n");
}

void ZmqRpcGenerator::GenerateStubPrologue(io::Printer &printer, const google::protobuf::FileDescriptor &file) const
{
    std::map<std::string, std::string> vars;
    vars["full_file_name"] = file.name();
    vars["file_name"] = fileName;
    std::string impl =
        "// Generated by the protocol buffer compiler.  DO NOT EDIT!\n"
        "// source: $full_file_name$\n"
        "#ifndef DATASYSTEM_PROTO_ZMQ_$file_name$_STUB_H\n"
        "#define DATASYSTEM_PROTO_ZMQ_$file_name$_STUB_H\n";
    if (file.message_type_count() > 0) {
        impl += "#include \"$file_name$.pb.h\"\n";
    }
    impl +=
        "#include \"datasystem/common/rpc/zmq/zmq_stub.h\"\n"
        "#include \"datasystem/common/rpc/zmq/zmq_message.h\"\n"
        "#include \"datasystem/common/util/net_util.h\"\n"
        "#include \"datasystem/utils/status.h\"\n"
        "#include \"datasystem/common/rpc/rpc_stub_base.h\"\n";
    printer.Print(vars, impl.c_str());
    for (auto k = 0; k < file.dependency_count(); ++k) {
        auto depend = file.dependency(k)->name();
        // Check if it ends in ".proto".
        bool isAProto = StripSuffix(depend, ".proto");
        if (isAProto) {
            vars["depend"] = depend;
            printer.Print(vars, "#include \"$depend$.pb.h\"\n");
        }
    }
    printer.Print(vars, "#include <string>\n");
}

void ZmqRpcGenerator::ImplementStubApiDecl(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc,
                                           const std::string &indent)
{
    for (auto j = 0; j < svc.method_count(); ++j) {
        if (svc.method(j) == nullptr) {
            continue;
        }
        auto &method = *(svc.method(j));
        if (method.client_streaming() && method.server_streaming()) {
            ImplementStubStreamingDecl(printer, method, indent);
        } else if (method.client_streaming() && !method.server_streaming()) {
            ImplementStubClientStreamingDecl(printer, method, indent);
        } else if (!method.client_streaming() && method.server_streaming()) {
            ImplementStubServerStreamingDecl(printer, method, indent);
        } else {
            ImplementStubAsyncWriteDecl(printer, method, indent);
            ImplementStubAsyncReadDecl(printer, method, indent);
            ImplementStubNoStreamDecl(printer, method, indent);
            ImplementStubNoStreamDecl2(printer, method, indent);
            ImplStubNoStreamShortDecl(printer, method, indent);
        }
    }
}

void ZmqRpcGenerator::ImplementGenericStubOtherFuncDecl(io::Printer &printer)
{
    const std::string otherFuncDecl =
        "    void ForgetRequest(int64_t tagId);\n"
        "    bool IsPeerAlive(uint32_t threshold);\n"
        "    void CacheSession(bool cache);\n"
        "    void SetExclusiveConnInfo(const std::optional<int32_t> &exclusiveId, const std::string &sockPath);\n";
    printer.PrintRaw(otherFuncDecl);
}

void ZmqRpcGenerator::ImplementStubAsyncWriteDecl(io::Printer &printer,
                                                  const google::protobuf::MethodDescriptor &method,
                                                  const std::string &indent)
{
    std::map<std::string, std::string> vars;
    vars["indent"] = indent;
    vars["method"] = method.name();
    vars["input_type"] = method.input_type()->name();
    vars["payload_send"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    printer.Print(vars,
                  "$indent$::datasystem::Status $method$AsyncWrite(const ::datasystem::RpcOptions &opt, "
                  "const $input_type$ &rq, int64_t &tagId$payload_send$);\n");
    printer.Print(
        vars,
        "$indent$::datasystem::Status $method$AsyncWrite(const $input_type$ &rq, int64_t &tagId$payload_send$);\n");
}

void ZmqRpcGenerator::ImplementStubStreamingDecl(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                 const std::string &indent)
{
    std::map<std::string, std::string> vars;
    vars["indent"] = indent;
    vars["method"] = method.name();
    vars["input_type"] = method.input_type()->name();
    vars["output_type"] = method.output_type()->name();
    printer.Print(vars,
                  "$indent$::datasystem::Status $method$(const ::datasystem::RpcOptions &opt, "
                  "std::unique_ptr<datasystem::ClientWriterReader<$input_type$, $output_type$>> *out);\n");
    printer.Print(vars,
                  "$indent$::datasystem::Status $method$"
                  "(std::unique_ptr<datasystem::ClientWriterReader<$input_type$, $output_type$>> *out);\n");
}

void ZmqRpcGenerator::ImplementStubClientStreamingDecl(io::Printer &printer,
                                                       const google::protobuf::MethodDescriptor &method,
                                                       const std::string &indent)
{
    std::map<std::string, std::string> vars;
    vars["indent"] = indent;
    vars["method"] = method.name();
    vars["input_type"] = method.input_type()->name();
    printer.Print(vars,
                  "$indent$::datasystem::Status $method$(const ::datasystem::RpcOptions &opt, "
                  "std::unique_ptr<datasystem::ClientWriter<$input_type$>> *out);\n");
    printer.Print(
        vars, "$indent$::datasystem::Status $method$(std::unique_ptr<datasystem::ClientWriter<$input_type$>> *out);\n");
}

void ZmqRpcGenerator::ImplementStubServerStreamingDecl(io::Printer &printer,
                                                       const google::protobuf::MethodDescriptor &method,
                                                       const std::string &indent)
{
    std::map<std::string, std::string> vars;
    vars["indent"] = indent;
    vars["method"] = method.name();
    vars["input_type"] = method.input_type()->name();
    vars["output_type"] = method.output_type()->name();
    vars["payload_send"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    printer.Print(
        vars,
        "$indent$::datasystem::Status $method$(const ::datasystem::RpcOptions &opt, "
        "std::unique_ptr<datasystem::ClientReader<$output_type$>> *out, const $input_type$ &rq$payload_send$);\n");
    printer.Print(
        vars,
        "$indent$::datasystem::Status $method$(std::unique_ptr<datasystem::ClientReader<$output_type$>> *out, "
        "const $input_type$ &rq$payload_send$);\n");
}

void ZmqRpcGenerator::ImplementStubAsyncReadDecl(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                 const std::string &indent)
{
    std::map<std::string, std::string> vars;
    vars["indent"] = indent;
    vars["method"] = method.name();
    vars["output_type"] = method.output_type()->name();
    vars["payload_recv"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";
    printer.Print(vars,
                  "$indent$::datasystem::Status $method$AsyncRead(int64_t tagId, $output_type$ &reply$payload_recv$, "
                  "::datasystem::RpcRecvFlags flags = ::datasystem::RpcRecvFlags::NONE);\n");
}

void ZmqRpcGenerator::ImplementStubNoStreamDecl(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                const std::string &indent)
{
    std::map<std::string, std::string> vars;
    vars["indent"] = indent;
    vars["method"] = method.name();
    vars["input_type"] = method.input_type()->name();
    vars["output_type"] = method.output_type()->name();
    vars["payload_send"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["payload_recv"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";
    printer.Print(vars,
                  "$indent$::datasystem::Status $method$(const ::datasystem::RpcOptions &opt, const $input_type$ &rq, "
                  "$output_type$ &reply$payload_send$$payload_recv$);\n");
}

void ZmqRpcGenerator::ImplementStubNoStreamDecl2(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                 const std::string &indent)
{
    std::map<std::string, std::string> vars;
    vars["indent"] = indent;
    vars["method"] = method.name();
    vars["input_type"] = method.input_type()->name();
    vars["output_type"] = method.output_type()->name();
    printer.Print(vars,
                  "$indent$::datasystem::Status $method$(const ::datasystem::RpcOptions &opt,\n"
                  "\t\tstd::unique_ptr<datasystem::ClientUnaryWriterReader<$input_type$, $output_type$>> *out);\n");
}

void ZmqRpcGenerator::ImplStubNoStreamShortDecl(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                const std::string &indent)
{
    std::map<std::string, std::string> vars;
    vars["indent"] = indent;
    vars["method"] = method.name();
    vars["input_type"] = method.input_type()->name();
    vars["output_type"] = method.output_type()->name();
    vars["payload_send"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["payload_recv"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";
    printer.Print(vars,
                  "$indent$::datasystem::Status $method$(const $input_type$ &rq, "
                  "$output_type$ &reply$payload_send$$payload_recv$);\n");
}

}  // namespace datasystem
